#include <iostream>
#include <stdio.h>
#include <opencv2/core/core.hpp>
#include <chrono>

#include "g2o/g2o/core/block_solver.h"
#include "g2o/g2o/core/optimization_algorithm_levenberg.h"
#include "g2o/g2o/solvers/linear_solver_eigen.h"
#include "g2o/g2o/types/types_six_dof_expmap.h"
#include "g2o/g2o/core/robust_kernel_impl.h"
#include "g2o/g2o/solvers/linear_solver_dense.h"
#include "g2o/g2o/types/types_seven_dof_expmap.h"

#include "Converter.h"


using namespace std;
using namespace cv;


static void monocularObservation(g2o::SparseOptimizer &optimizer, vector<g2o::EdgeSE3ProjectXYZOnlyPose*> &vpEdgesMono,
					vector<size_t> &vnIndexEdgeMono, int number_i, float deltaMono, float invSigma,
					float keypoint_x, float keypoint_y, float fx, float fy, float cx, float cy, float x, float y, float z)
{
	Eigen::Matrix<double,2,1> obs;
    obs << keypoint_x, keypoint_y;

    g2o::EdgeSE3ProjectXYZOnlyPose* e = new g2o::EdgeSE3ProjectXYZOnlyPose();

    e->setVertex(0, dynamic_cast<g2o::OptimizableGraph::Vertex*>(optimizer.vertex(0)));
    e->setMeasurement(obs);
    const float invSigma2 = invSigma;
    e->setInformation(Eigen::Matrix2d::Identity()*invSigma2);

    g2o::RobustKernelHuber* rk = new g2o::RobustKernelHuber;
    e->setRobustKernel(rk);
    rk->setDelta(deltaMono);

    e->fx = fx;
    e->fy = fy;
    e->cx = cx;
    e->cy = cy;
    e->Xw[0] = x;
 	e->Xw[1] = y;
    e->Xw[2] = z;

    optimizer.addEdge(e);

    vpEdgesMono.push_back(e);
    vnIndexEdgeMono.push_back(number_i);
}



static void stereoObservation(g2o::SparseOptimizer &optimizer, vector<g2o::EdgeStereoSE3ProjectXYZOnlyPose*> &vpEdgesStereo,
	vector<size_t> &vnIndexEdgeStereo, int number_i, float deltaStereo, float invSigma,
	float keypoint_x, float keypoint_y, float keypoint_ur, float fx, float fy, float cx, float cy, float baselineFx,
	float x, float y, float z)
{
	//SET EDGE
	Eigen::Matrix<double,3,1> obs;
	
	obs << keypoint_x, keypoint_y, keypoint_ur;
	
	g2o::EdgeStereoSE3ProjectXYZOnlyPose* e = new g2o::EdgeStereoSE3ProjectXYZOnlyPose();
	
	e->setVertex(0, dynamic_cast<g2o::OptimizableGraph::Vertex*>(optimizer.vertex(0)));
	e->setMeasurement(obs);
	const float invSigma2 = invSigma;
	Eigen::Matrix3d Info = Eigen::Matrix3d::Identity()*invSigma2;
	e->setInformation(Info);
	
	g2o::RobustKernelHuber* rk = new g2o::RobustKernelHuber;
	e->setRobustKernel(rk);
	rk->setDelta(deltaStereo);
	
	e->fx = fx;
	e->fy = fy;
	e->cx = cx;
	e->cy = cy;
	e->bf = baselineFx;

	e->Xw[0] = x;
	e->Xw[1] = y;
	e->Xw[2] = z;
	
	optimizer.addEdge(e);
	
	vpEdgesStereo.push_back(e);
	vnIndexEdgeStereo.push_back(number_i);

}


int main()
{
	std::chrono::steady_clock::time_point t1 = std::chrono::steady_clock::now();
	g2o::SparseOptimizer optimizer;
    g2o::BlockSolver_6_3::LinearSolverType * linearSolver;

    linearSolver = new g2o::LinearSolverDense<g2o::BlockSolver_6_3::PoseMatrixType>();

    g2o::BlockSolver_6_3 * solver_ptr = new g2o::BlockSolver_6_3(linearSolver);

    g2o::OptimizationAlgorithmLevenberg* solver = new g2o::OptimizationAlgorithmLevenberg(solver_ptr);
    optimizer.setAlgorithm(solver);

    int nInitialCorrespondences=0;

	cv::Mat RTcw = (Mat_<double>(4,4) << 	1.0, 0., 0., 0.,
											0., 1.0, 0., 0.,
											0., 0., 1.0, 0.,
											0., 0., 0., 1.0);

    // Set Frame vertex
    g2o::VertexSE3Expmap * vSE3 = new g2o::VertexSE3Expmap();
    vSE3->setEstimate(Converter::toSE3Quat(RTcw));
    vSE3->setId(0);
    vSE3->setFixed(false);
    optimizer.addVertex(vSE3);

    // Set MapPoint vertices
    int N = 1200;

    vector<g2o::EdgeSE3ProjectXYZOnlyPose*> vpEdgesMono;
    vector<size_t> vnIndexEdgeMono;
    vpEdgesMono.reserve(N);
    vnIndexEdgeMono.reserve(N);

    vector<g2o::EdgeStereoSE3ProjectXYZOnlyPose*> vpEdgesStereo;
    vector<size_t> vnIndexEdgeStereo;
    vpEdgesStereo.reserve(N);
    vnIndexEdgeStereo.reserve(N);

    const float deltaMono = sqrt(5.991);
    const float deltaStereo = sqrt(7.815);

	cout << "poseOptimization test" << endl;

	std::vector<bool> m_vectorOutlier;
	m_vectorOutlier = vector<bool>(N, false);
	

	//i=6
	monocularObservation(optimizer, vpEdgesMono, vnIndexEdgeMono, 6, deltaMono, 1.0,
						640, 326, 
						435.204681, 435.204681, 367.451721, 252.200851,
						2.284418, 0.618562, 3.647755);


	//i= 10
	stereoObservation(optimizer, vpEdgesStereo, vnIndexEdgeStereo, 10, deltaStereo, 1.0,
					227, 247, 198.583725,
					435.204681, 435.204681, 367.451721, 252.200851, 47.906395,
					-0.541096, -0.020037, 1.676645);

	monocularObservation(optimizer, vpEdgesMono, vnIndexEdgeMono, 12, deltaMono, 1.0,
						600, 271, 
						435.204681, 435.204681, 367.451721, 252.200851,
						0.646260, 0.049464, 1.209449);

	monocularObservation(optimizer, vpEdgesMono, vnIndexEdgeMono, 13, deltaMono, 1.0,
						595, 270, 
						435.204681, 435.204681, 367.451721, 252.200851,
						1.949803, 0.152516, 3.729159);

	stereoObservation(optimizer, vpEdgesStereo, vnIndexEdgeStereo, 17, deltaStereo, 1.0,
					652, 238, 639.238831,
					435.204681, 435.204681, 367.451721, 252.200851, 47.906395,
					2.474312, -0.123485, 3.784357);

	stereoObservation(optimizer, vpEdgesStereo, vnIndexEdgeStereo, 18, deltaStereo, 1.0,
					648, 239, 635.036804,
					435.204681, 435.204681, 367.451721, 252.200851, 47.906395,
					2.393394, -0.112618, 3.712789);

	stereoObservation(optimizer, vpEdgesStereo, vnIndexEdgeStereo, 22, deltaStereo, 1.0,
					726, 275, 712.329285,
					435.204681, 435.204681, 367.451721, 252.200851, 47.906395,
					2.882585, 0.183296, 3.498872);

	stereoObservation(optimizer, vpEdgesStereo, vnIndexEdgeStereo, 27, deltaStereo, 1.0,
					678.000000, 278.000000, 667.146057,
					435.204681, 435.204681, 367.451721, 252.200851, 47.906395,
					3.243993, 0.269498, 4.546156);

	monocularObservation(optimizer, vpEdgesMono, vnIndexEdgeMono, 38, deltaMono, 1.0,
						422, 292, 
						435.204681, 435.204681, 367.451721, 252.200851,
						0.374998, 0.278713, 3.047735);

	stereoObservation(optimizer, vpEdgesStereo, vnIndexEdgeStereo, 42, deltaStereo, 1.0,
					679.000000, 106.000000, 665.805481,
					435.204681, 435.204681, 367.451721, 252.200851, 47.906395,
					2.598130, -1.219230, 3.629353);

	stereoObservation(optimizer, vpEdgesStereo, vnIndexEdgeStereo, 43, deltaStereo, 1.0,
					446.000000, 121.000000, 433.123108,
					435.204681, 435.204681, 367.451721, 252.200851, 47.906395,
					0.669836, -1.118842, 3.711297);

	stereoObservation(optimizer, vpEdgesStereo, vnIndexEdgeStereo, 45, deltaStereo, 1.0,
					446.000000, 114.000000, 433.795044,
					435.204681, 435.204681, 367.451721, 252.200851, 47.906395,
					0.708694, -1.246904, 3.926593);

	stereoObservation(optimizer, vpEdgesStereo, vnIndexEdgeStereo, 51, deltaStereo, 1.0,
					131.000000, 271.000000, 97.151611,
					435.204681, 435.204681, 367.451721, 252.200851, 47.906395,
					-0.767002, 0.060981, 1.411718);

	monocularObservation(optimizer, vpEdgesMono, vnIndexEdgeMono, 54, deltaMono, 1.0,
						226, 274, 
						435.204681, 435.204681, 367.451721, 252.200851,
						-0.471067, 0.072596, 1.449333);

	stereoObservation(optimizer, vpEdgesStereo, vnIndexEdgeStereo, 68, deltaStereo, 1.0,
					301.000000, 265.000000, 279.984406,
					435.204681, 435.204681, 367.451721, 252.200851, 47.906395,
					-0.339838, 0.070570, 2.225664);

	stereoObservation(optimizer, vpEdgesStereo, vnIndexEdgeStereo, 71, deltaStereo, 1.0,
					357.000000, 270.000000, 344.053192,
					435.204681, 435.204681, 367.451721, 252.200851, 47.906395,
					-0.088746, 0.151133, 3.695344);

	stereoObservation(optimizer, vpEdgesStereo, vnIndexEdgeStereo, 75, deltaStereo, 1.0,
					682.000000, 429.000000, 652.803528,
					435.204681, 435.204681, 367.451721, 252.200851, 47.906395,
					1.190776, 0.669303, 1.647541);

	stereoObservation(optimizer, vpEdgesStereo, vnIndexEdgeStereo, 76, deltaStereo, 1.0,
					566.000000, 242.000000, 553.002563,
					435.204681, 435.204681, 367.451721, 252.200851, 47.906395,
					1.672270, -0.085917, 3.665506);

	stereoObservation(optimizer, vpEdgesStereo, vnIndexEdgeStereo, 86, deltaStereo, 1.0,
					448.000000, 180.000000, 434.535339,
					435.204681, 435.204681, 367.451721, 252.200851, 47.906395,
					0.660885, -0.592395, 3.570779);

	stereoObservation(optimizer, vpEdgesStereo, vnIndexEdgeStereo, 87, deltaStereo, 1.0,
					533.000000, 237.000000, 520.221436,
					435.204681, 435.204681, 367.451721, 252.200851, 47.906395,
					1.424286, -0.139383, 3.744262);

	stereoObservation(optimizer, vpEdgesStereo, vnIndexEdgeStereo, 88, deltaStereo, 1.0,
					240.000000, 298.000000, 210.993790,
					435.204681, 435.204681, 367.451721, 252.200851, 47.906395,
					-0.481886, 0.173163, 1.645477);

	stereoObservation(optimizer, vpEdgesStereo, vnIndexEdgeStereo, 94, deltaStereo, 1.0,
					244.000000, 159.000000, 231.805466,
					435.204681, 435.204681, 367.451721, 252.200851, 47.906395,
					-1.109530, -0.837648, 3.911428);

	stereoObservation(optimizer, vpEdgesStereo, vnIndexEdgeStereo, 110, deltaStereo, 1.0,
					576.000000, 223.000000, 562.038940,
					435.204681, 435.204681, 367.451721, 252.200851, 47.906395,
					1.632720, -0.228613, 3.407209);

	stereoObservation(optimizer, vpEdgesStereo, vnIndexEdgeStereo, 118, deltaStereo, 1.0,
					655.000000, 332.000000, 640.426575,
					435.204681, 435.204681, 367.451721, 252.200851, 47.906395,
					2.152742, 0.597420, 3.258178);

	stereoObservation(optimizer, vpEdgesStereo, vnIndexEdgeStereo, 120, deltaStereo, 1.0,
					601.000000, 234.000000, 587.766174,
					435.204681, 435.204681, 367.451721, 252.200851, 47.906395,
					1.939290, -0.151132, 3.613763);

	


	cout << "start optimization" << endl;

	//start optimization
	// We perform 4 optimizations, after each optimization we classify observation as inlier/outlier
    // At the next optimization, outliers are not included, but at the end they can be classified as inliers again.
    const float chi2Mono[4]={5.991,5.991,5.991,5.991};
    const float chi2Stereo[4]={7.815,7.815,7.815, 7.815};
    const int its[4]={10,10,10,10};    

    int nBad=0;
    for(size_t it=0; it<4; it++)
    {

        vSE3->setEstimate(Converter::toSE3Quat(RTcw));
        optimizer.initializeOptimization(0);
        optimizer.optimize(its[it]);

        nBad=0;
        for(size_t i=0, iend=vpEdgesMono.size(); i<iend; i++)
        {
            g2o::EdgeSE3ProjectXYZOnlyPose* e = vpEdgesMono[i];

            const size_t idx = vnIndexEdgeMono[i];

            if(m_vectorOutlier[idx])
            {
                e->computeError();
            }

            const float chi2 = e->chi2();

            if(chi2>chi2Mono[it])
            {                
                m_vectorOutlier[idx]=true;
                e->setLevel(1);
                nBad++;
            }
            else
            {
                m_vectorOutlier[idx]=false;
                e->setLevel(0);
            }

            if(it==2)
                e->setRobustKernel(0);
        }

        for(size_t i=0, iend=vpEdgesStereo.size(); i<iend; i++)
        {
            g2o::EdgeStereoSE3ProjectXYZOnlyPose* e = vpEdgesStereo[i];

            const size_t idx = vnIndexEdgeStereo[i];

            if(m_vectorOutlier[idx])
            {
                e->computeError();
            }

            const float chi2 = e->chi2();

            if(chi2>chi2Stereo[it])
            {
                m_vectorOutlier[idx]=true;
                e->setLevel(1);
                nBad++;
            }
            else
            {                
                e->setLevel(0);
                m_vectorOutlier[idx]=false;
            }

            if(it==2)
                e->setRobustKernel(0);
        }

        if(optimizer.edges().size()<10)
            break;
    }    

    // Recover optimized pose and return number of inliers
    g2o::VertexSE3Expmap* vSE3_recov = static_cast<g2o::VertexSE3Expmap*>(optimizer.vertex(0));
    g2o::SE3Quat SE3quat_recov = vSE3_recov->estimate();
    cv::Mat pose = Converter::toCvMat(SE3quat_recov);

	std::chrono::steady_clock::time_point t2 = std::chrono::steady_clock::now();
	double ttrack = std::chrono::duration_cast<std::chrono::duration<double>>(t2 - t1).count();
	cout << "PoseOptimization--------ttrack:" << ttrack * 1e3 << " ms" << endl;
	return 0;
}

